package inference;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;

import com.analog.lyric.dimple.model.Discrete;

import model.ModelParameters;
import model.PredicateInfo;
import model.TableInformation;
import utils.IntPair;
import utils.Pair;

public class PredicateFactor {

	
	public static Pair<int[][], double[]> createFactor(Discrete parent_type_variable, Discrete predicate_variable,
													   ModelParameters params) {
		
		Object[] predicates = predicate_variable.getDomain().getElements();
		int unk_pred_id = predicates.length - 1;
		Object[] parent_types = parent_type_variable.getDomain().getElements();
		int unk_type_id = parent_types.length - 1;
		
		// Set up output data structures.
		int init_size = parent_types.length + predicates.length - 1;
		ArrayList<IntPair> indices = new ArrayList<IntPair>(init_size);
		ArrayList<Double> probabilities = new ArrayList<Double>(init_size);
		
		HashMap<Integer, Integer> type_to_cnt = new HashMap<Integer, Integer>();
		
		for (int pred = 0; pred < predicates.length; pred++) {
			if (pred == unk_pred_id) {
				continue;
			}
			Integer predicate_id = (Integer)(predicates[pred]);
			PredicateInfo pred_info = params.getTableInformation().getPredicateInfo(predicate_id);
			HashSet<Integer> cooccurring_types = params.getTableInformation().getCooccurringTypes(pred_info.getDomainId());
			
			for (int par_type = 0; par_type < parent_types.length; par_type++) {
				Integer par_type_id = (Integer)(parent_types[par_type]);
				if (cooccurring_types.contains(par_type_id)) {
					Integer cnt = type_to_cnt.get(par_type);
					if (cnt == null) {
						type_to_cnt.put(par_type, 1);
					} else {
						type_to_cnt.put(par_type, cnt.intValue() + 1);
					}
				}
			}
		}
		
		// P(pred | unk) = uniform
		for (int pred = 0; pred < predicates.length; pred++) {
			indices.add(new IntPair(unk_type_id, pred));
			probabilities.add(1.0 / predicates.length);
		}
		
		for (int par_type = 0; par_type < parent_types.length; par_type++) {
			if (par_type == unk_type_id) {
				continue;
			}
			
			Integer cnt_int = type_to_cnt.get(par_type);
			int cnt = 1;
			if (cnt_int != null)
				cnt = cnt_int.intValue() + 1;
			indices.add(new IntPair(par_type, unk_pred_id));
			if (cnt == 1)
				probabilities.add(1.0);
			else
				probabilities.add((1.0 / cnt) * (1.0 / parent_types.length));
		}
		
		for (int pred = 0; pred < predicates.length; pred++) {
			if (pred == unk_pred_id) {
				continue;
			}
			Integer predicate_id = (Integer)(predicates[pred]);
			PredicateInfo pred_info = params.getTableInformation().getPredicateInfo(predicate_id);
			HashSet<Integer> cooccurring_types = params.getTableInformation().getCooccurringTypes(pred_info.getDomainId());
			
			for (int par_type = 0; par_type < parent_types.length; par_type++) {
				if (par_type == unk_type_id) {
					continue;
				}
				Integer par_type_id = (Integer)(parent_types[par_type]);

				if (!cooccurring_types.contains(par_type_id)) {
					continue;
				}
				
				int cnt = type_to_cnt.get(par_type).intValue() + 1;
				double prob = (1.0 / cnt) + ((1.0 / cnt) * (((double)parent_types.length - 1) / parent_types.length) * (1.0 / (cnt - 1)));
				indices.add(new IntPair(par_type, pred));
				probabilities.add(prob);
			}
		}
		
		return translateListsToPrimitives(indices, probabilities);
	}
	

	
	private static Pair<int[][], double[]> translateListsToPrimitives(ArrayList<IntPair> indices, ArrayList<Double> probabilities) {
		int[][] indices_arr = new int[indices.size()][2];
		double[] prob_arr = new double[probabilities.size()];
		
		Iterator<IntPair> it = indices.iterator();
		Iterator<Double> it_prob = probabilities.iterator();
		for (int i = 0; i < indices.size(); i++) {
			IntPair next = it.next();
			indices_arr[i] = next.toArray();
			prob_arr[i] = it_prob.next().doubleValue();
//			System.out.println(next.toString() + " " + prob_arr[i]);
		}
		
		return new Pair<int[][], double[]>(indices_arr, prob_arr);
	}
	
}